import json
import xml.etree.ElementTree as ET
from loguru import logger
import os # Import os
import uuid # Import uuid
import subprocess # Import subprocess

class Utils:
    def __init__(self, root_agent_instance):
        self.root_agent = root_agent_instance

    def create_error_response(self, task_id: str, error_message: str, status: str = "failed", progress: int = 0) -> dict:
        """Creates a standardized dictionary for error responses."""
        self.root_agent.task_states[task_id] = {
            "status": status,
            "progress": progress,
            "result": "",
            "error": error_message,
        }
        return {
            "task_id": task_id,
            "final_result": f"Error: {error_message}",
            "llm_usage_metrics": self.root_agent.llm_usage_metrics,
        }

    def parse_params(self, params: str) -> dict | str:
        """Parses the params string, handling XML, CDATA, and plain strings."""
        logger.debug(f"_parse_params received: {params!r} (type: {type(params)})")

        if isinstance(params, dict):
            return params
        
        if not isinstance(params, str) or not params.strip():
            logger.debug("_parse_params: params is empty or not a string, returning empty dict.")
            return {} # Return empty dict for empty or non-string params

        try:
            # Attempt to parse params as XML
            parsed_params_element = ET.fromstring(params)
            final_params = {}
            logger.debug(f"_parse_params: Successfully parsed XML. Root tag: {parsed_params_element.tag}")

            # Check if it's a CDATA section (for execute_in_sandbox)
            if parsed_params_element.text and parsed_params_element.text.strip().startswith('<code_execution>'):
                cdata_content = parsed_params_element.text.strip()
                logger.debug(f"_parse_params: Detected CDATA content: {cdata_content!r}")
                # Remove the outer CDATA markers if they are still present
                if cdata_content.startswith('<![CDATA[') and cdata_content.endswith(']]>'):
                    cdata_content = cdata_content[9:-3].strip()
                
                # Remove the outer curly braces if they are still present
                if cdata_content.startswith('{') and cdata_content.endswith('}'):
                    cdata_content = cdata_content[1:-1].strip()

                logger.debug(f"_parse_params: CDATA content after stripping: {cdata_content!r}")
                cdata_root = ET.fromstring(cdata_content)
                for child in cdata_root:
                    final_params[child.tag] = child.text
            # If it's a single value wrapped in <value>
            elif parsed_params_element.tag == "value":
                final_params = parsed_params_element.text
            else:
                for child in parsed_params_element:
                    final_params[child.tag] = child.text
            logger.debug(f"_parse_params: Parsed XML params to dict/value: {final_params!r} (type: {type(final_params)})")
            logger.debug(f"_parse_params: Returning final_params: {final_params!r} (type: {type(final_params)})")
            return final_params
        except ET.ParseError as e:
            logger.debug(f"_parse_params: Params is not valid XML ({e}), treating as plain string: {params!r}")
            # If not an XML string, use params as is
            return params
        except Exception as e:
            logger.error(f"_parse_params: Unexpected error during params parsing: {e}. Raw params: {params!r}", exc_info=True)
            raise Exception(f"Error processing params: {e}") # Re-raise as generic Exception for now

    def run_python_code(self, code: str, timeout: int = 60) -> dict:
        """
        Executes Python code in a subprocess and captures its output.
        Args:
            code (str): The Python code to execute.
            timeout (int): Maximum time in seconds to wait for the code to execute.
        Returns:
            dict: A dictionary containing 'stdout', 'stderr', and 'error' (if any).
        """
        try:
            # Create a temporary file to write the code to
            temp_file_path = os.path.join(self.root_agent.base_work_dir, "temp", f"temp_code_{uuid.uuid4().hex}.py")
            os.makedirs(os.path.dirname(temp_file_path), exist_ok=True)
            with open(temp_file_path, "w", encoding="utf-8") as f:
                f.write(code)

            process = subprocess.run(
                ["python", temp_file_path],
                capture_output=True,
                text=True,
                check=False, # Do not raise an exception for non-zero exit codes
                timeout=timeout
            )
            
            # Clean up the temporary file
            os.remove(temp_file_path)

            return {
                "stdout": process.stdout.strip(),
                "stderr": process.stderr.strip(),
                "error": None if process.returncode == 0 else f"Process exited with code {process.returncode}"
            }
        except subprocess.TimeoutExpired:
            # Clean up the temporary file even if timeout occurs
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            return {
                "stdout": "",
                "stderr": f"Execution timed out after {timeout} seconds.",
                "error": "Timeout"
            }
        except Exception as e:
            # Clean up the temporary file even if other exceptions occur
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            return {
                "stdout": "",
                "stderr": str(e),
                "error": "Exception during execution"
            }
